package org.example.shardUpload;

import org.example.nettyrpc.EndpointBase;
import org.example.nettyrpc.common.IpPortTuple;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.ArrayList;
import java.util.HashSet;

import static org.example.shardUpload.common.UploadProperties.File_Shard_SIZE;

public class UploadClient extends EndpointBase {

    public UploadClient(IpPortTuple serverIpPort) {
        super(null, new IpPortTuple[]{serverIpPort});
    }

    public String getMD5checksum(String filename) throws IOException {
        InputStream fileInStream = Files.newInputStream(new File(filename).toPath());
        MessageDigest messageDigest;
        try {
            messageDigest = MessageDigest.getInstance("MD5");
        } catch (NoSuchAlgorithmException ex) {
            throw new RuntimeException("NoSuchAlgorithmException");
        }

        int shardNum = (fileInStream.available() + File_Shard_SIZE - 1) / File_Shard_SIZE;
        byte[] buf = new byte[File_Shard_SIZE];
        int shardLen = fileInStream.read(buf, 0, File_Shard_SIZE);

        for (int i = 0; i < shardNum; i++) {
            messageDigest.update(buf, 0, shardLen);
        }

        byte[] digest = messageDigest.digest();
        StringBuilder sb = new StringBuilder();
        for (byte b : digest) {
            sb.append(String.format("%02x", b));
        }

        fileInStream.close();
        return sb.toString();
    }

    public void uploadFileByName(String filename) throws IOException, InterruptedException {
        String md5checksum = getMD5checksum(filename);
        InputStream fileInStream = Files.newInputStream(new File(filename).toPath());
        int shardNum = (fileInStream.available() + File_Shard_SIZE - 1) / File_Shard_SIZE;
        for (; ; ) {
            ArrayList<Integer> sent = (ArrayList<Integer>) remoteCall(0, "searchSentSuccessShard", filename, md5checksum);
            HashSet<Integer> sentSet = new HashSet<>(sent);

            if (sentSet.size() == shardNum) {
                // then merge shards
                Object mergeResult = remoteCall(0, "requestMergeShards", filename, shardNum, md5checksum);
                if (mergeResult instanceof String) {
                    throw new RuntimeException((String) mergeResult);
                }
                boolean mergeSuccess = (boolean) mergeResult;
                if (mergeSuccess) {
                    return;
                } else {
                    continue;
                }
            }

            for (int i = 0; i < shardNum; i++) {
                if (!sentSet.contains(i)) {
                    sendshards(fileInStream, i, filename, md5checksum);
                }
            }
        }
    }

    private void sendshards(InputStream fileInStream, int shardIndex, String filename, String md5checksum) throws InterruptedException, IOException {
        byte[] buf = new byte[File_Shard_SIZE];
        int shardLen = fileInStream.read(buf, 0, File_Shard_SIZE);
        remoteCall(0, "requestUpLoadShard", filename, shardIndex, shardLen, buf, md5checksum);
    }

}
